#!/usr/bin/env python3
import rospy
import serial
import struct
import threading
import numpy as np
from geometry_msgs.msg import PoseStamped

class WirelessTransceiver:
    def __init__(self):
        rospy.init_node('uav_wireless_transceiver_uav1')
        
        # 配置参数
        self.uav_id = 0x01  # 本机ID
        self.serial_port = rospy.get_param('~port', '/dev/ttyUSB0')
        self.baudrate = rospy.get_param('~baud', 115200)
        
        # 目标地址配置
        self.ground_station_addr = bytes([0x00, 0x04, 0x17])
        self.uav_addrs = {
            0x02: bytes([0x00, 0x02, 0x17]),  # UAV2
            0x03: bytes([0x00, 0x03, 0x17])   # UAV3
        }
        
        # 预编译数据格式（提升打包效率）
        self.pose_dtype = np.dtype([
            ('pos', np.float16, 3),      # x,y,z
            ('orient', np.float16, 4),   # x,y,z,w
            ('padding', np.float16)      # 填充位
        ])
        
        # 串口初始化
        try:
            self.ser = serial.Serial(
                self.serial_port,
                self.baudrate,
                timeout=0.01,          # 缩短超时
                write_timeout=0.01,    # 防止写阻塞
                xonxoff=False,         # 禁用流控
                rtscts=False            # 禁用硬件流控
            )
            rospy.loginfo(f"Serial port {self.serial_port} opened successfully")
        except serial.SerialException as e:
            rospy.logerr(f"Failed to open serial port: {e}")
            exit(1)
            
        # 订阅本机数据
        rospy.Subscriber('/mavros/local_position/pose', PoseStamped, self.mavros_pose_cb)
        rospy.Subscriber('/detect/local_position/pose', PoseStamped, self.detect_pose_cb)
        
        # 存储最新数据
        self.current_mavros_pose = None
        self.current_detect_pose = None
        
        # 接收缓冲区
        self.rx_buffer = bytearray()
        
        # 启动线程
        self.running = True
        self.tx_thread = threading.Thread(target=self.tx_worker)
        self.rx_thread = threading.Thread(target=self.rx_worker)
        self.tx_thread.start()
        self.rx_thread.start()
        
        rospy.on_shutdown(self.shutdown)
    
    def shutdown(self):
        self.running = False
        self.tx_thread.join()
        self.rx_thread.join()
        self.ser.close()
        rospy.loginfo("Wireless transceiver shutdown")
    
    def mavros_pose_cb(self, msg):
        self.current_mavros_pose = msg
    
    def detect_pose_cb(self, msg):
        self.current_detect_pose = msg
    
    def calculate_checksum(self, data):
        return sum(data) % 256
    
    def pack_pose_data(self):
        """使用半精度浮点数打包位置数据（原56字节 → 现28字节）"""
        if self.current_mavros_pose is None:
            return None
            
        # 创建结构化数组
        pose_array = np.zeros(2, dtype=self.pose_dtype)
        
        # 填充MAVROS数据
        pose_array[0]['pos'] = [
            self.current_mavros_pose.pose.position.x,
            self.current_mavros_pose.pose.position.y,
            self.current_mavros_pose.pose.position.z
        ]
        pose_array[0]['orient'] = [
            self.current_mavros_pose.pose.orientation.x,
            self.current_mavros_pose.pose.orientation.y,
            self.current_mavros_pose.pose.orientation.z,
            self.current_mavros_pose.pose.orientation.w
        ]
        
        # 填充Detect数据（或默认值）
        detect_pose = self.current_detect_pose if self.current_detect_pose else PoseStamped()
        pose_array[1]['pos'] = [
            detect_pose.pose.position.x,
            detect_pose.pose.position.y,
            detect_pose.pose.position.z
        ]
        pose_array[1]['orient'] = [
            detect_pose.pose.orientation.x,
            detect_pose.pose.orientation.y,
            detect_pose.pose.orientation.z,
            detect_pose.pose.orientation.w if detect_pose.pose.orientation.w != 0 else 1.0
        ]
        
        return pose_array.tobytes()
    
    def tx_worker(self):
        """发送线程（时间精确控制版本）"""
        ground_interval = 1.5  # 5Hz
        uav_interval = 0.8     # 10Hz
        next_ground_tx = rospy.Time.now().to_sec()
        next_uav_tx = rospy.Time.now().to_sec()

        while self.running:
            now = rospy.Time.now().to_sec()
            
            # 地面站发送（严格5Hz）
            if now >= next_ground_tx:
                if self.current_mavros_pose:
                    self.send_data(self.ground_station_addr)
                next_ground_tx += ground_interval
            
            # 无人机间发送（严格10Hz）
            if now >= next_uav_tx:
                if self.current_mavros_pose:
                    for addr in self.uav_addrs.values():
                        self.send_data(addr)
                next_uav_tx += uav_interval
            
            # 动态休眠（剩余时间的50%作为安全余量）
            sleep_time = min(
                max(0, next_ground_tx - now) * 0.5,
                max(0, next_uav_tx - now) * 0.5
            )
            rospy.sleep(max(0.001, sleep_time))
    
    def send_data(self, target_addr):
        """优化后的数据发送方法"""
        pose_data = self.pack_pose_data()
        if not pose_data:
            return
            
        # 协议帧结构
        header = bytes([0xAA, 0x55])
        footer = bytes([0x55, 0xAA])
        uav_id = bytes([self.uav_id])
        
        # 计算校验和（包含头）
        checksum_data = header + uav_id + pose_data
        checksum = bytes([self.calculate_checksum(checksum_data)])
        
        # 完整数据包（目标地址3 + 头2 + ID1 + 数据28 + 校验1 + 尾2 = 37字节）
        packet = target_addr + header + uav_id + pose_data + checksum + footer
        
        try:
            self.ser.write(packet)
            rospy.logdebug(f"Sent {len(packet)} bytes to {target_addr.hex()}")
        except serial.SerialException as e:
            rospy.logwarn(f"Send error: {e}")
    
    def rx_worker(self):
        """接收线程（适配半精度数据）"""
        header = bytes([0xAA, 0x55])
        footer = bytes([0x55, 0xAA])
        min_frame_len = 2 + 1 + 28 + 1 + 2  # 34字节
        
        while self.running:
            try:
                # 非阻塞读取
                data = self.ser.read(self.ser.in_waiting or 1)
                if data:
                    rospy.loginfo("Raw RX Data: " + " ".join([f"{byte:02X}" for byte in data]))
                    self.rx_buffer += data
            except serial.SerialException as e:
                rospy.logwarn(f"Serial error: {e}")
                continue
            
            # 帧处理
            while len(self.rx_buffer) >= min_frame_len:
                # 查找帧头
                header_pos = self.rx_buffer.find(header)
                if header_pos < 0:
                    self.rx_buffer.clear()
                    break
                
                # 移除无效数据
                if header_pos > 0:
                    self.rx_buffer = self.rx_buffer[header_pos:]
                
                # 检查完整帧
                if len(self.rx_buffer) < min_frame_len:
                    break
                
                # 查找帧尾
                footer_pos = self.rx_buffer.find(footer, 2)
                if footer_pos < 0:
                    break
                
                # 提取并解析帧
                frame = self.rx_buffer[:footer_pos+2]
                self.rx_buffer = self.rx_buffer[footer_pos+2:]
                self.parse_frame(frame)
            
            rospy.sleep(0.001)
    
    def parse_frame(self, frame):
        """解析半精度数据帧"""
        try:
            # 基本验证
            if len(frame) < 34:
                return
            if frame[0] != 0xAA or frame[1] != 0x55:
                return
            if frame[-2] != 0x55 or frame[-1] != 0xAA:
                return
            
            # 提取字段
            sender_id = frame[2]
            data = frame[3:31]
            received_checksum = frame[31]
            
            # 校验和验证
            if self.calculate_checksum(frame[:31]) != received_checksum:
                return
            
            # 解析数据
            pose_array = np.frombuffer(data, dtype=self.pose_dtype)
            
            # 转换为PoseStamped
            mavros_pose = PoseStamped()
            mavros_pose.pose.position.x = float(pose_array[0]['pos'][0])
            mavros_pose.pose.position.y = float(pose_array[0]['pos'][1])
            mavros_pose.pose.position.z = float(pose_array[0]['pos'][2])
            mavros_pose.pose.orientation.x = float(pose_array[0]['orient'][0])
            mavros_pose.pose.orientation.y = float(pose_array[0]['orient'][1])
            mavros_pose.pose.orientation.z = float(pose_array[0]['orient'][2])
            mavros_pose.pose.orientation.w = float(pose_array[0]['orient'][3])
            
            detect_pose = PoseStamped()
            detect_pose.pose.position.x = float(pose_array[1]['pos'][0])
            detect_pose.pose.position.y = float(pose_array[1]['pos'][1])
            detect_pose.pose.position.z = float(pose_array[1]['pos'][2])
            detect_pose.pose.orientation.x = float(pose_array[1]['orient'][0])
            detect_pose.pose.orientation.y = float(pose_array[1]['orient'][1])
            detect_pose.pose.orientation.z = float(pose_array[1]['orient'][2])
            detect_pose.pose.orientation.w = float(pose_array[1]['orient'][3])
            
            # 发布数据
            self.publish_received_data(sender_id, mavros_pose, detect_pose)
            
        except Exception as e:
            rospy.logerr(f"Parse error: {e}")
    
    def publish_received_data(self, sender_id, mavros_pose, detect_pose):
        """发布接收到的数据"""
        if sender_id not in self.uav_addrs:
            return
        
        # 动态创建发布者
        mavros_topic = f'/uav{sender_id}/mavros/local_position/pose'
        detect_topic = f'/uav{sender_id}/detect/local_position/pose'
        
        if not hasattr(self, f'uav{sender_id}_mavros_pub'):
            setattr(self, f'uav{sender_id}_mavros_pub', 
                   rospy.Publisher(mavros_topic, PoseStamped, queue_size=1))
            setattr(self, f'uav{sender_id}_detect_pub',
                   rospy.Publisher(detect_topic, PoseStamped, queue_size=1))
        
        # 设置时间戳并发布
        stamp = rospy.Time.now()
        
        mavros_pose.header.stamp = stamp
        mavros_pose.header.frame_id = f'uav{sender_id}_local'
        getattr(self, f'uav{sender_id}_mavros_pub').publish(mavros_pose)
        
        detect_pose.header.stamp = stamp
        detect_pose.header.frame_id = f'uav{sender_id}_detect'
        getattr(self, f'uav{sender_id}_detect_pub').publish(detect_pose)

if __name__ == '__main__':
    try:
        transceiver = WirelessTransceiver()
        rospy.spin()
    except rospy.ROSInterruptException:
        pass